package org.example.petitionplatformsystem.TongyiChat;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import io.reactivex.Flowable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class AiWebSocketHandler extends TextWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(AiWebSocketHandler.class);

    private final Generation generation;

    @Value("${ai.api.key}")
    private String appKey;

    @Value("${ai.model}")
    private String model;

    private final Map<Long, List<Message>> userConversationHistory = new ConcurrentHashMap<>();
    private final Map<Long, Integer> userLastSentIndex = new ConcurrentHashMap<>();

    @Autowired
    public AiWebSocketHandler(Generation generation) {
        this.generation = generation;
        if (this.generation == null) {
            logger.error("Generation service is not initialized in constructor.");
        } else {
            logger.info("Generation service successfully injected in constructor.");
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        try {
            Long userId = Long.parseLong(session.getUri().getPath().split("/")[3]);
            String question = message.getPayload();

            if (generation == null) {
                logger.error("Generation service is not initialized.");
                session.sendMessage(new TextMessage("Generation service is not available."));
                return;
            }

            if (appKey == null || appKey.isEmpty()) {
                logger.error("API key is not set.");
                session.sendMessage(new TextMessage("API key is not set."));
                return;
            }

            logger.info("Received message from user {}: {}", userId, question);
            List<Message> conversationHistory = userConversationHistory.computeIfAbsent(userId, k -> new ArrayList<>());
            final int[] lastSentIndex = {userLastSentIndex.getOrDefault(userId, 0)};

            conversationHistory.add(Message.builder().role(Role.USER.getValue()).content(question).build());

            if (conversationHistory.size() > 6) {
                conversationHistory = conversationHistory.subList(conversationHistory.size() - 6, conversationHistory.size());
            }

            GenerationParam generationParam = GenerationParam.builder()
                    .model(model)
                    .messages(conversationHistory)
                    .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                    .topP(0.8)
                    .apiKey(appKey)
                    .build();

            Flowable<GenerationResult> result = generation.streamCall(generationParam);

            List<Message> finalConversationHistory = conversationHistory;
            Flux.from(result)
                    .flatMap(m -> {
                        String newContent = m.getOutput().getChoices().get(0).getMessage().getContent();
                        String incrementalContent = newContent.substring(lastSentIndex[0]);
                        lastSentIndex[0] = newContent.length();
                        userLastSentIndex.put(userId, lastSentIndex[0]);

                        finalConversationHistory.add(Message.builder().role(Role.ASSISTANT.getValue()).content(incrementalContent).build());
                        return Mono.just(incrementalContent);
                    })
                    .doOnComplete(() -> {
                        lastSentIndex[0] = 0;
                        userLastSentIndex.remove(userId);
                    })
                    .subscribeOn(Schedulers.boundedElastic())
                    .subscribe(content -> {
                        try {
                            session.sendMessage(new TextMessage(content));
                        } catch (Exception e) {
                            logger.error("Error sending message to user {}: {}", userId, e.getMessage());
                        }
                    });
        } catch (NumberFormatException e) {
            logger.error("Invalid user ID format: {}", e.getMessage());
            session.sendMessage(new TextMessage("Invalid user ID format."));
        } catch (Exception e) {
            logger.error("Error processing message: {}", e.getMessage());
            session.sendMessage(new TextMessage("Error processing message."));
        }
    }
}
